#include <iostream>
#include <vector>
#include <queue>
#include <thread>
#include <mutex>
#include <condition_variable>
#include <future>
#include <functional>
#include <stdexcept>
#include <type_traits>
#include <concepts>
#include <chrono>
#include <optional>

using namespace std::chrono_literals;

class ThreadPool {
public:
    explicit ThreadPool(size_t num_threads = std::thread::hardware_concurrency())
        : stop_source_{} 
    {
        workers_.reserve(num_threads);
        for(size_t i = 0; i < num_threads; ++i) {
            workers_.emplace_back([this](std::stop_token st) {
                worker_thread(st);
            });
        }
    }

    ~ThreadPool() {
        shutdown();
    }

    // 提交任务到线程池
    template<typename F, typename... Args>
    requires std::invocable<F, Args...>
    auto enqueue(F&& f, Args&&... args) 
        -> std::optional<std::future<std::invoke_result_t<F, Args...>>>
    {
        using return_type = std::invoke_result_t<F, Args...>;
        
        // 检查是否已关闭
        if (stop_source_.stop_requested()) {
            return std::nullopt;
        }
        
        auto task = std::make_shared<std::packaged_task<return_type()>>(
            [func = std::forward<F>(f), ...args = std::forward<Args>(args)]() mutable {
                return func(std::forward<Args>(args)...);
            }
        );
        
        std::future<return_type> result = task->get_future();
        
        {
            std::unique_lock lock(queue_mutex_);
            tasks_.emplace([task]() { (*task)(); });
        }
        
        cv_.notify_one();
        return result;
    }

    // 关闭线程池（优雅停止）
    void shutdown() noexcept {
        if (stop_source_.stop_requested()) return;
        
        // 请求停止
        stop_source_.request_stop();
        
        // 唤醒所有线程
        cv_.notify_all();
        
        // 清空任务队列
        clear_tasks();
        
        // 等待所有线程结束
        workers_.clear(); // jthread 析构时自动 join
    }

    // 等待所有任务完成
    void wait_completion() {
        std::unique_lock lock(queue_mutex_);
        completion_cv_.wait(lock, [this] {
            return tasks_.empty() && active_tasks_ == 0;
        });
    }

    // 检查是否已停止
    bool is_shutdown() const noexcept {
        return stop_source_.stop_requested();
    }

    // 获取活动线程数
    size_t active_threads() const noexcept {
        return active_tasks_;
    }

    // 获取等待任务数
    size_t pending_tasks() {
        std::lock_guard lock(queue_mutex_);
        return tasks_.size();
    }

    // 获取线程池大小
    size_t size() const noexcept {
        return workers_.size();
    }

private:
    void worker_thread(std::stop_token st) {
        while (!st.stop_requested()) {
            std::function<void()> task;
            
            // 获取任务
            {
                std::unique_lock lock(queue_mutex_);
                
                // 等待任务或停止请求
                cv_.wait(lock, st, [this] { 
                    return !tasks_.empty() || stop_source_.stop_requested();
                });
                
                // 检查停止请求
                if (stop_source_.stop_requested()) {
                    break;
                }
                
                // 获取任务
                task = std::move(tasks_.front());
                tasks_.pop();
                
                // 更新活动任务计数
                active_tasks_++;
            }
            
            // 执行任务
            try {
                task();
            } catch (const std::exception& e) {
                std::cerr << "ThreadPool task exception: " << e.what() << "\n";
            } catch (...) {
                std::cerr << "ThreadPool task unknown exception\n";
            }
            
            // 更新任务计数
            {
                std::lock_guard lock(queue_mutex_);
                active_tasks_--;
                
                // 检查是否所有任务都已完成
                if (tasks_.empty() && active_tasks_ == 0) {
                    completion_cv_.notify_all();
                }
            }
        }
    }

    void clear_tasks() noexcept {
        std::lock_guard lock(queue_mutex_);
        while (!tasks_.empty()) {
            tasks_.pop();
        }
    }

    // 成员声明顺序至关重要
    std::queue<std::function<void()>> tasks_;         // 任务队列
    std::mutex queue_mutex_;                          // 保护任务队列的互斥量
    std::condition_variable_any cv_;                  // 任务可用通知
    std::condition_variable completion_cv_;           // 任务完成通知
    std::stop_source stop_source_;                    // 停止控制
    
    // 状态变量
    size_t active_tasks_ = 0;                         // 当前正在执行的任务数
    
    // 工作线程最后声明 - 确保最先析构
    std::vector<std::jthread> workers_;               // 工作线程
};

// 示例使用
int main() {
    try {
        ThreadPool pool(4);
        
        std::cout << "ThreadPool created with " << pool.size() << " threads\n";
        
        // 提交任务
        std::vector<std::future<int>> results;
        for (int i = 0; i < 10; ++i) {
            auto result = pool.enqueue([i] {
                std::this_thread::sleep_for(100ms);
                std::cout << "Task " << i << " executed by thread " 
                          << std::this_thread::get_id() << "\n";
                return i * i;
            });
            
            if (result) {
                results.push_back(std::move(*result));
            }
        }
        
        // 等待所有任务完成
        pool.wait_completion();
        std::cout << "\nAll tasks completed!\n";
        
        // 获取结果
        for (auto& fut : results) {
            std::cout << "Result: " << fut.get() << "\n";
        }
        
        // 提交更多任务
        for (int i = 10; i < 15; ++i) {
            pool.enqueue([i] {
                std::this_thread::sleep_for(50ms);
                std::cout << "Additional task " << i << "\n";
            });
        }
        
        // 等待 300ms 让任务开始执行
        std::this_thread::sleep_for(300ms);
        
        // 关闭线程池
        std::cout << "\nShutting down thread pool...\n";
        pool.shutdown();
        
        // 尝试提交新任务（应该失败）
        auto result = pool.enqueue([] {
            std::cout << "This should not be executed!\n";
            return 42;
        });
        
        if (!result) {
            std::cout << "Task submission failed (pool is shutdown)\n";
        }
        
        std::cout << "ThreadPool shutdown complete\n";
        
    } catch (const std::exception& e) {
        std::cerr << "Error: " << e.what() << "\n";
        return 1;
    }
    
    return 0;
}
